import asyncio
import os
import json
import sys
import logging
from typing import Optional, List, Dict, Any
from contextlib import AsyncExitStack

from openai import OpenAI  
from dotenv import load_dotenv

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

# 加载 .env 文件，确保 API Key 受到保护
load_dotenv()

# 配置日志，将信息输出到控制台和文件
logger = logging.getLogger("MCPClientLogger")
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter("[%(asctime)s] %(levelname)s: %(message)s")

# 控制台日志处理器
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.DEBUG)
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)

# 文件日志处理器，日志保存到 mcp_client_trace.txt
file_handler = logging.FileHandler("mcp_client_trace.txt", mode="w", encoding="utf-8")
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)


class MCPClient:
    def __init__(self):
        """初始化 MCP 客户端，并设置日志记录"""
        self.exit_stack = AsyncExitStack()
        self.openai_api_key = os.getenv("OPENAI_API_KEY")  # 读取 OpenAI API Key
        self.base_url = os.getenv("BASE_URL")  # 读取 BASE URL
        self.model = os.getenv("MODEL")  # 读取模型
        if not self.openai_api_key:
            error_msg = "❌ 未找到 OpenAI API Key，请在 .env 文件中设置 OPENAI_API_KEY"
            logger.error(error_msg)
            raise ValueError(error_msg)
        self.client = OpenAI(api_key=self.openai_api_key, base_url=self.base_url)  # 创建 OpenAI client

        # 保存多个服务器的会话和工具映射
        self.sessions: List[ClientSession] = []
        self.tool_session_map: Dict[str, ClientSession] = {}  # 工具名 -> 会话
        self.tool_details_map: Dict[str, Dict[str, Any]] = {}   # 工具名 -> 工具详细信息

    async def connect_to_servers(self, server_script_paths: List[str]) -> List[Dict[str, Any]]:
        """
        连接到多个 MCP 服务器并汇总各个服务器的工具列表，同时记录每一步日志信息。
        :param server_script_paths: 各服务器脚本的路径列表
        :return: 所有服务器合并的工具列表
        """
        for path in server_script_paths:
            logger.info(f"准备连接服务器脚本 {path}")
            # 判断服务器脚本类型
            is_python = path.endswith('.py')
            is_js = path.endswith('.js')
            if not (is_python or is_js):
                err_msg = "服务器脚本必须是 .py 或 .js 文件"
                logger.error(err_msg)
                raise ValueError(err_msg)

            command = "python" if is_python else "node"
            server_params = StdioServerParameters(
                command=command,
                args=[path],
                env=None
            )
            logger.debug(f"启动 MCP 服务器 {path} 使用命令 {command}")
            # 启动 MCP 服务器并建立通信
            stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
            stdio, write = stdio_transport
            session = await self.exit_stack.enter_async_context(ClientSession(stdio, write))
            await session.initialize()

            # 获取当前服务器上提供的工具
            logger.info(f"请求服务器 {path} 的工具列表")
            response = await session.list_tools()
            tools = response.tools
            logger.info(f"已连接到服务器 {path}，支持的工具: {[tool.name for tool in tools]}")

            # 保存工具对应的会话和详细信息
            for tool in tools:
                self.tool_session_map[tool.name] = session
                self.tool_details_map[tool.name] = {
                    "name": tool.name,
                    "description": tool.description,
                    "input_schema": tool.inputSchema # 工具的输入参数结构
                }
                logger.debug(f"记录工具：{tool.name}，描述：{tool.description}, 输入参数结构：{tool.inputSchema}")
            self.sessions.append(session)

        # 组合所有工具，生成给大模型使用的工具描述列表
        available_tools = []
        for details in self.tool_details_map.values():
            available_tools.append({
                "type": "function",
                "function": {
                    "name": details["name"],
                    "description": details["description"],
                    "input_schema": details["input_schema"]
                }
            })
        logger.info("所有服务器的工具列表已整合完成")
        return available_tools

    async def process_query(self, query: str) -> str:
        """
        使用大模型处理查询并调用合并后的 MCP 工具 (Function Calling)。
        在调用每一步时都记录详细的结构信息，包括发送给大模型的请求结构、
        大模型返回的工具调用请求结构等，方便调试和查看。
        """
        logger.info(f"收到用户查询: {query}")
        messages = [{"role": "user", "content": query}]
        
        # 构造可用工具列表
        available_tools = []
        for details in self.tool_details_map.values():
            available_tools.append({
                "type": "function",
                "function": {
                    "name": details["name"],
                    "description": details["description"],
                    "input_schema": details["input_schema"]
                }
            })
        logger.debug(f"构造的工具列表: {json.dumps(available_tools, ensure_ascii=False, indent=2)}")
        
        # 记录发送给大模型的初始请求结构
        request_payload = {
            "model": self.model,
            "messages": messages,
            "tools": available_tools
        }
        logger.info(f"发送给大模型的初始请求结构: {json.dumps(request_payload, ensure_ascii=False, indent=2)}")
        
        # 调用大模型生成回复
        logger.info("调用 OpenAI 大模型生成回复")
        response = self.client.chat.completions.create(
            model=self.model,            
            messages=messages,
            tools=available_tools     
        )
        
        # 记录大模型返回的完整响应结构
        logger.info(f"大模型返回的完整响应: {json.dumps(response.model_dump(), ensure_ascii=False, indent=2)}")
        
        content = response.choices[0]
        if content.finish_reason == "tool_calls":
            # 当模型要求调用工具时，解析工具调用相关参数
            tool_call = content.message.tool_calls[0]
            tool_name = tool_call.function.name
            tool_args = json.loads(tool_call.function.arguments)
            logger.info(f"大模型请求调用工具 {tool_name}，传入参数: {json.dumps(tool_args, ensure_ascii=False, indent=2)}")
            
            # 根据工具名称查找对应的会话
            session = self.tool_session_map.get(tool_name)
            if not session:
                err_msg = f"未找到工具 {tool_name} 对应的服务器会话"
                logger.error(err_msg)
                raise ValueError(err_msg)
            
            # 调用对应工具
            logger.info(f"调用工具 {tool_name} 并传入参数 {json.dumps(tool_args, ensure_ascii=False)}")
            result = await session.call_tool(tool_name, tool_args)
            logger.info(f"工具 {tool_name} 返回: {result.content[0].text}")

            # 将工具调用记录和返回结果记录到对话记录中
            messages.append(content.message.model_dump())
            messages.append({
                "role": "tool",
                "content": result.content[0].text,
                "tool_call_id": tool_call.id,
            })
            logger.debug(f"更新后的对话记录: {json.dumps(messages, ensure_ascii=False, indent=2)}")
            
            # 记录发送给大模型的请求结构 (包含工具调用记录)
            final_request_payload = {
                "model": self.model,
                "messages": messages,
            }
            logger.info(f"发送给大模型的最终请求结构: {json.dumps(final_request_payload, ensure_ascii=False, indent=2)}")
            
            # 第二次调用大模型，生成最终回答
            logger.info("调用大模型生成最终回答")
            final_response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
            )
            # 记录大模型最终返回的响应结构
            logger.info(f"大模型最终返回的完整响应: {json.dumps(final_response.model_dump(), ensure_ascii=False, indent=2)}")
            final_content = final_response.choices[0].message.content
            logger.info(f"大模型的最终回答: {final_content}")
            return final_content
            
        logger.info("大模型返回直接回答，无工具调用")
        return content.message.content

    async def chat_loop(self):
        """运行交互式聊天循环，并记录整个对话过程"""
        logger.info("🤖 MCP 客户端已启动！")
        print("\n🤖 MCP 客户端已启动！输入 'quit' 退出")
        while True:
            try:
                query = input("\n你: ").strip()
                logger.info(f"用户输入: {query}")
                if query.lower() == 'quit':
                    logger.info("用户退出聊天")
                    break
                
                response = await self.process_query(query)
                logger.info(f"回复结果: {response}")
                print(f"\n🤖 OpenAI: {response}")

            except Exception as e:
                err_msg = f"发生错误: {str(e)}"
                logger.error(err_msg)
                print(f"\n⚠️ {err_msg}")

    async def cleanup(self):
        """清理资源，在退出时捕获 ProcessLookupError 异常，同时记录清理日志"""
        try:
            logger.info("开始清理资源")
            await self.exit_stack.aclose()
            logger.info("资源清理完成")
        except ProcessLookupError:
            logger.warning("子进程已退出，忽略 ProcessLookupError")


async def main():
    if len(sys.argv) < 2:
        usage_msg = "Usage: python client_test.py <path_to_server_script1> [<path_to_server_script2> ...]"
        logger.error(usage_msg)
        print(usage_msg)
        sys.exit(1)

    client = MCPClient()
    try:
        logger.info("连接所有传入的服务器脚本")
        await client.connect_to_servers(sys.argv[1:])
        await client.chat_loop()
    finally:
        logger.info("程序结束，开始资源清理")
        await client.cleanup()

if __name__ == "__main__":
    asyncio.run(main())
